from pydantic import BaseModel, Field
from typing import Dict, Optional, List
from enum import Enum

class GenerationType(str, Enum):
    """代码生成类型枚举"""
    AUTH = "AUTH"
    ACCESS_CONTROL = "ACCESS_CONTROL"
    KEY_DATA_TRANSPORT_CONFIDENTIALITY_AND_INTEGRITY = "KEY_DATA_TRANSPORT_CONFIDENTIALITY_AND_INTEGRITY"
    KEY_DATA_STORE_CONFIDENTIALITY = "KEY_DATA_STORE_CONFIDENTIALITY"
    KEY_DATA_STORE_INTEGRITY = "KEY_DATA_STORE_INTEGRITY"
    KEY_DATA_SAFE_MARK_INTEGRITY = "KEY_DATA_SAFE_MARK_INTEGRITY"
    NON_REPUDIATION = "NON-REPUDIATION"

class ProgrammingLanguage(str, Enum):
    """编程语言枚举"""
    PYTHON = "python"
    JAVA = "java"
    JAVASCRIPT = "javascript"
    CSHARP = "csharp"
    GO = "go"

class CodeGenerationRequest(BaseModel):
    """代码生成请求模型"""
    key_data: str = Field(..., description="关键数据类型，如'KEY'")
    gen_type: GenerationType = Field(..., description="代码生成类型")
    crypto_eval_extract: Dict[str, str] = Field(..., description="提取的密码测评信息")
    gen_language: ProgrammingLanguage = Field(default=ProgrammingLanguage.PYTHON, description="目标编程语言")

class CodeGenerationResponse(BaseModel):
    """代码生成响应模型"""
    success: bool = Field(..., description="是否成功")
    code: Optional[str] = Field(None, description="生成的代码")
    message: str = Field(..., description="响应消息")
    error: Optional[str] = Field(None, description="错误信息")

class DocumentProcessRequest(BaseModel):
    """文档处理请求模型"""
    doc_path: str = Field(..., description="文档路径")

class DocumentProcessResponse(BaseModel):
    """文档处理响应模型"""
    success: bool = Field(..., description="是否成功")
    crypto_eval_information: Optional[Dict[str, str]] = Field(None, description="原始文档信息")
    crypto_eval_extract: Optional[Dict[str, str]] = Field(None, description="结构化提取信息")
    message: str = Field(..., description="响应消息")
    error: Optional[str] = Field(None, description="错误信息")

class APIInfoResponse(BaseModel):
    """API信息响应模型"""
    success: bool = Field(..., description="是否成功")
    api_types: List[str] = Field(..., description="支持的API类型列表")
    message: str = Field(..., description="响应消息")

class SupportedTypesResponse(BaseModel):
    """支持类型响应模型"""
    success: bool = Field(..., description="是否成功")
    generation_types: List[str] = Field(..., description="支持的代码生成类型")
    programming_languages: List[str] = Field(..., description="支持的编程语言")
    message: str = Field(..., description="响应消息")

class HealthCheckResponse(BaseModel):
    """健康检查响应模型"""
    status: str = Field(..., description="服务状态")
    version: str = Field(..., description="服务版本")
    timestamp: str = Field(..., description="检查时间") 